#
# tcifar10_02.py
#
############################################################
#
from PIL import Image
import matplotlib.pyplot as plt  # plt 用于显示图片
import numpy as np
import torch
import torch.nn as nn
import torchvision
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import torch.utils.data as data

from utils import (ROOT)
from models import ConvNet

#
# 1. basic setting
#
# 设备设置
# torch.cuda.set_device(0)  # 这句用来设置pytorch在哪块GPU上运行，pytorch-cpu版本不需要运行这句
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# 超参数设置
num_epochs = 5
num_classes = 10    # 分类数量
batch_size = 32
learning_rate = 0.001
#
# 2. Load CIFAR-10 datasets and use Data Augmentation
#
# cifar10 分类索引
classes = ('plane', 'car', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck')

# 数据增广方法
transform = transforms.Compose([
    # +4填充至36x36
    transforms.Pad(4),
    # 随机水平翻转
    transforms.RandomHorizontalFlip(),
    # 随机裁剪至32x32
    transforms.RandomCrop(32),
    # 转换至Tensor
    transforms.ToTensor(),
])

# CIFAR-10 dataset
train_dataset = datasets.CIFAR10(
    root=ROOT, train=True, transform=transform, download=True)
test_dataset = datasets.CIFAR10(
    root=ROOT, train=False, transform=transform)

# Data loader
train_loader = data.DataLoader(
    dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = data.DataLoader(
    dataset=test_dataset, batch_size=batch_size, shuffle=False)

#
# 3. View an image and its corresponding label
#
# 查看数据, 取一组batch
data_iter = iter(test_loader)
images, labels = next(data_iter)

# 取batch中的一张图像
idx = 15
image = images[idx].numpy()             # image: (3, 32, 32)
image = np.transpose(image, (1, 2, 0))  # image: (32, 32, 3)
# plt.imshow(image)
# plt.savefig("expamle.png")
print(classes[labels[idx].numpy()])     # ship

#
# Design the net and train、eval and save the model
#
# 实例化一个模型，并迁移至gpu
model = ConvNet(num_classes).to(device)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# model.train()
total_step = len(train_loader)
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        # 注意模型在GPU中，数据也要搬到GPU中
        images = torch.Tensor(images).to(device)
        labels = torch.Tensor(labels).to(device)

        # 前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)

        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (i + 1) % 100 == 0:
            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch +
                  1, num_epochs, i + 1, total_step, loss.item()))


# 设置为评估模式
model.eval()
# 节省计算资源，不去计算梯度
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    print('Test Accuracy of the model on the test images: {} %'.format(
        100 * correct / total))


#  保存模型
# torch.save(model.state_dict(), 'model.ckpt')

# 测试集的图像
# 查看数据,取一组batch
data_iter = iter(test_loader)
images, labels = next(data_iter)

# 取测试集图像测试
# 取batch中的一张图像，显示图像和真实标签
idx = 10
image = images[idx].numpy()
image = np.transpose(image, (1, 2, 0))
plt.imshow(image)
plt.savefig("example2.png")
print(classes[labels[idx].numpy()])


# 转换为（B,C,H,W）大小
imagebatch = image.reshape(-1, 3, 32, 32)
# 转换为torch tensor
image_tensor = torch.from_numpy(imagebatch)

# 调用模型进行评估
model.eval()
output = model(image_tensor.to(device))
_, predicted = torch.max(output.data, 1)
pre = predicted.cpu().numpy()
print(pre)   # 查看预测结果ID
print(classes[pre[0]])

# 
# 自定义图像测试
#
# 读入一张自己图像进行测试
srcPath = 'horse.jpg'
image_temporary = np.array(plt.imread(srcPath))
newImg = np.array(Image.fromarray(
    image_temporary).resize((32, 32)))  # (3, 32, 32)

plt.imshow(newImg)  # 显示图片
plt.axis('off')  # 不显示坐标轴
plt.savefig("horse_resize.png")

# 转换为（B,C,H,W）大小
imagebatch = newImg.reshape(-1, 3, 32, 32)
# 转换为torch tensor
image_tensor = torch.from_numpy(imagebatch).float()
# 调用模型进行评估
model.eval()
output = model(image_tensor.to(device))
_, predicted = torch.max(output.data, 1)
pre = predicted.cpu().numpy()
print(pre)  # 查看预测结果ID
print(classes[pre[0]])
